######################### 将markdown格式的内容进行切片，根据标题进行切片，支持处理超长文本 #########################

def chunk_markdown(markdown_content, max_chunk_size=None, min_chunk_size=None):
    """将markdown格式的内容进行切片，根据标题进行切片，支持处理超长文本
    
    Args:
        markdown_content: 要处理的markdown文本
        max_chunk_size: 最大块大小（字符数），如果一个块超过此大小，会尝试进一步分割
        min_chunk_size: 最小块大小（字符数），小于此大小的块会尝试合并
    
    Returns:
        list: 切分后的文本块列表
    """
    import re
    from itertools import chain

    # 如果内容为空，直接返回空列表
    if not markdown_content or not markdown_content.strip():
        return []
    
    # 标准化换行符，确保使用LF
    markdown_content = markdown_content.replace('\r\n', '\n')
    
    # 使用更宽松的正则表达式匹配标题
    # 标题是行首的#号，后面可以有也可以没有空格，再跟文本
    pattern = re.compile(r'^(#+)(\s*)(.*)$', re.MULTILINE)
    
    # 查找所有标题并保存其位置和级别
    headers = []
    for match in pattern.finditer(markdown_content):
        # 忽略空标题（只有#号没有内容的行）
        if not match.group(3).strip():
            continue
        headers.append((
            match.start(),                   # 标题在文本中的起始位置
            match.group(0),                  # 完整标题文本
            len(match.group(1))              # 标题级别(#的数量)
        ))
    
    # 如果没有标题，将整个文本作为一个块
    if not headers:
        if markdown_content.strip():
            return [markdown_content.strip()]
        return []
    
    # 处理文档开头没有标题的情况
    initial_chunks = []
    if headers[0][0] > 0:
        content_before_first_header = markdown_content[:headers[0][0]].strip()
        if content_before_first_header:
            initial_chunks.append(content_before_first_header)
    
    # 基于标题提取块
    section_chunks = []
    for i, (pos, header, level) in enumerate(headers):
        # 计算当前标题内容的范围
        start = pos
        end = headers[i+1][0] if i < len(headers) - 1 else len(markdown_content)
        
        # 提取内容
        section_text = markdown_content[start:end].strip()
        
        # 如果设置了最大块大小且当前块超过限制
        if max_chunk_size and len(section_text) > max_chunk_size:
            # 尝试在段落处进一步分割超长块
            paragraph_chunks = _split_by_paragraphs(section_text, max_chunk_size)
            section_chunks.extend(paragraph_chunks)
        else:
            section_chunks.append(section_text)
    
    # 合并所有块
    all_chunks = list(chain(initial_chunks, section_chunks))
    
    # 如果设置了最小块大小，尝试合并小块
    if min_chunk_size:
        all_chunks = _merge_small_chunks(all_chunks, min_chunk_size)
    
    return all_chunks

def _split_by_paragraphs(text, max_size):
    """根据段落分割超长文本块
    
    Args:
        text: 要分割的文本
        max_size: 最大块大小
        
    Returns:
        list: 分割后的文本块列表
    """
    import re
    
    # 分割成段落
    paragraphs = re.split(r'\n\s*\n', text)
    
    chunks = []
    current_chunk = []
    current_size = 0
    
    # 标题行（如果有）
    header = None
    if paragraphs and re.match(r'^#+\s+', paragraphs[0]):
        header = paragraphs[0]
        paragraphs = paragraphs[1:]
    
    for para in paragraphs:
        para_size = len(para)
        
        # 如果当前段落就已经超过最大大小，强制分割这个段落
        if para_size > max_size:
            # 如果当前块不为空，先保存
            if current_chunk:
                chunks.append('\n\n'.join(current_chunk))
                current_chunk = []
                current_size = 0
            
            # 分割超长段落（按句子或简单地按字符数）
            if header:
                para_chunks = _split_large_paragraph(para, max_size - len(header) - 2)
                # 为每个分割后的块添加原标题
                chunks.extend([f"{header}\n\n{chunk}" for chunk in para_chunks])
            else:
                chunks.extend(_split_large_paragraph(para, max_size))
            
            continue
        
        # 如果添加此段落会超过限制，先保存当前块
        if current_size + para_size + (2 if current_chunk else 0) > max_size:
            if current_chunk:
                chunks.append('\n\n'.join(current_chunk))
                current_chunk = []
                current_size = 0
        
        # 添加段落到当前块
        if header and not current_chunk:
            current_chunk.append(header)
            current_size += len(header)
            if para:
                current_chunk.append(para)
                current_size += para_size + 2  # +2 for '\n\n'
        else:
            current_chunk.append(para)
            current_size += para_size + (2 if len(current_chunk) > 1 else 0)  # +2 for '\n\n'
    
    # 保存最后一个块
    if current_chunk:
        chunks.append('\n\n'.join(current_chunk))
    
    return chunks

def _split_large_paragraph(paragraph, max_size):
    """分割超大段落
    
    Args:
        paragraph: 要分割的段落
        max_size: 最大块大小
        
    Returns:
        list: 分割后的段落列表
    """
    import re
    
    # 尝试按句子分割
    sentences = re.split(r'(?<=[.!?])\s+', paragraph)
    
    chunks = []
    current_chunk = []
    current_size = 0
    
    for sentence in sentences:
        sentence_size = len(sentence)
        
        # 如果单个句子就超过了最大大小，按字符分割
        if sentence_size > max_size:
            # 先保存当前块
            if current_chunk:
                chunks.append(' '.join(current_chunk))
            
            # 简单地按字符数分割
            for i in range(0, sentence_size, max_size):
                chunks.append(sentence[i:i + max_size])
            
            current_chunk = []
            current_size = 0
            continue
        
        # 如果添加此句子会超过限制，先保存当前块
        if current_size + sentence_size + (1 if current_chunk else 0) > max_size:
            chunks.append(' '.join(current_chunk))
            current_chunk = []
            current_size = 0
        
        # 添加句子到当前块
        current_chunk.append(sentence)
        current_size += sentence_size + (1 if len(current_chunk) > 1 else 0)  # +1 for space
    
    # 保存最后一个块
    if current_chunk:
        chunks.append(' '.join(current_chunk))
    
    return chunks

def _merge_small_chunks(chunks, min_size):
    """合并小块
    
    Args:
        chunks: 文本块列表
        min_size: 最小块大小
        
    Returns:
        list: 合并后的文本块列表
    """
    if not chunks:
        return []
    
    import re
    
    result = []
    current_chunk = chunks[0]
    current_size = len(current_chunk)
    
    for i in range(1, len(chunks)):
        next_chunk = chunks[i]
        next_size = len(next_chunk)
        
        # 如果当前块小于最小大小，且下一个块不是以标题开始，尝试合并
        if current_size < min_size and not re.match(r'^#+\s+', next_chunk):
            current_chunk = f"{current_chunk}\n\n{next_chunk}"
            current_size += next_size + 2  # +2 for '\n\n'
        else:
            result.append(current_chunk)
            current_chunk = next_chunk
            current_size = next_size
    
    # 添加最后一个块
    result.append(current_chunk)
    
    return result

######################### 文本embedding #########################

import numpy as np
from zhipuai import ZhipuAI
import os
import time
from typing import List, Union, Optional

def get_zhipuai_embedding(text: Union[str, List[str]], model: str = "embedding-2", 
                          retry_count: int = 3, retry_delay: float = 1.0) -> Union[np.ndarray, List[np.ndarray]]:
    """
    使用智谱AI的embedding模型将文本转换为向量
    
    Args:
        text: 需要转换的文本，可以是单个字符串或字符串列表
        model: embedding模型名称，可选"embedding-2"或"embedding-3"
        retry_count: 请求失败时的重试次数
        retry_delay: 重试之间的延迟时间(秒)
    
    Returns:
        如果输入是单个字符串，返回对应的向量；如果是字符串列表，返回向量列表
    
    Raises:
        ValueError: 文本长度超过最大限制或发送空文本
        TypeError: 输入类型错误
        Exception: API调用失败
    """
    client = ZhipuAI(api_key=os.environ.get("ZHIPUAI_API_KEY", "9e2458a8acf46274d1c1b9418eec500c.oBojpibVHt3IyXnz"))
    
    # 根据模型设置限制参数
    if model == "embedding-2":
        # embedding-2: 单条最多512 tokens，数组总长不超过8K
        single_token_limit = 512
        total_token_limit = 8000
        max_array_size = float('inf')  # 没有条数限制，只有总token限制
    elif model == "embedding-3":
        # embedding-3: 单条最多3072 tokens，数组最大不超过64条
        single_token_limit = 3072
        total_token_limit = 3072 * 64  # 理论上限
        max_array_size = 64
    else:
        raise ValueError(f"不支持的模型: {model}，请选择 'embedding-2' 或 'embedding-3'")
    
    # 验证并处理输入
    if isinstance(text, str):
        if not text.strip():
            raise ValueError("输入文本不能为空")
        inputs = [text]
        is_single = True
    elif isinstance(text, list) and all(isinstance(item, str) for item in text):
        if not text or any(not item.strip() for item in text):
            raise ValueError("输入文本列表不能包含空文本")
        if len(text) > max_array_size:
            raise ValueError(f"模型 {model} 一次最多支持 {max_array_size} 条文本")
        inputs = text
        is_single = False
    else:
        raise TypeError("输入必须是字符串或字符串列表")
    
    # 截断过长文本以符合token限制
    # 这里使用简单的字符截断，实际应用中应使用分词器计算真实token数
    truncated_inputs = []
    total_tokens = 0
    
    for input_text in inputs:
        # 简单估算：中文约1字符1token，英文约4字符1token
        approx_tokens = len(input_text)
        
        # 检查单条文本是否超过限制
        if approx_tokens > single_token_limit:
            # 截断文本并添加警告
            truncated_text = input_text[:single_token_limit]
            print(f"警告：文本超过 {single_token_limit} tokens限制，已自动截断")
            truncated_inputs.append(truncated_text)
            total_tokens += single_token_limit
        else:
            truncated_inputs.append(input_text)
            total_tokens += approx_tokens
    
    # 检查总token数是否超过限制
    if total_tokens > total_token_limit:
        raise ValueError(f"总token数 {total_tokens} 超过模型 {model} 的限制 {total_token_limit}")
    
    # 带重试的API调用
    for attempt in range(retry_count):
        try:
            response = client.embeddings.create(
                model=model,
                input=truncated_inputs
            )
            
            # 处理返回结果
            embeddings = [np.array(item.embedding) for item in response.data]
            return embeddings[0] if is_single else embeddings
            
        except Exception as e:
            # 最后一次尝试失败，抛出异常
            if attempt == retry_count - 1:
                raise Exception(f"获取embedding失败: {str(e)}")
            
            # 其他尝试失败，等待后重试
            print(f"API调用失败，{retry_delay}秒后重试: {str(e)}")
            time.sleep(retry_delay)
            # 增加退避时间以应对速率限制
            retry_delay *= 1.5

def batch_get_zhipuai_embedding(texts: List[str], model: str = "embedding-2", 
                               rate_limit_pause: float = 0.5) -> List[np.ndarray]:
    """
    批量获取文本向量，处理大量文本时避免API限制
    
    Args:
        texts: 文本列表
        model: embedding模型名称，可选"embedding-2"或"embedding-3"
        rate_limit_pause: 每批次之间的暂停时间(秒)，避免触发速率限制
    
    Returns:
        文本向量列表
    """
    if not texts:
        return []
    
    # 根据模型设置批处理参数
    if model == "embedding-2":
        # embedding-2: 单条最多512 tokens，数组总长不超过8K
        single_token_limit = 512
        max_total_tokens = 8000
        batch_size = min(len(texts), max_total_tokens // single_token_limit)  # 估算批次大小
    elif model == "embedding-3":
        # embedding-3: 单条最多3072 tokens，数组最大不超过64条
        single_token_limit = 3072
        batch_size = 64  # 每批最多64条
    else:
        raise ValueError(f"不支持的模型: {model}，请选择 'embedding-2' 或 'embedding-3'")
    
    all_embeddings = []
    
    # 按批次处理
    for i in range(0, len(texts), batch_size):
        # 获取当前批次
        batch_texts = texts[i:i + batch_size]
        
        try:
            # 处理当前批次
            batch_embeddings = get_zhipuai_embedding(batch_texts, model)
            all_embeddings.extend(batch_embeddings)
            
            # 暂停以避免速率限制
            if i + batch_size < len(texts):
                time.sleep(rate_limit_pause)
                
        except Exception as e:
            print(f"批次处理失败 ({i} 到 {i + batch_size}): {str(e)}")
            # 对于embedding-2，如果批次太大，尝试减少批次大小再处理
            if model == "embedding-2" and "token" in str(e).lower() and batch_size > 1:
                # 减小批次大小，逐个处理
                print("尝试减小批次大小并重新处理...")
                for j in range(i, min(i + batch_size, len(texts))):
                    try:
                        single_embedding = get_zhipuai_embedding(texts[j], model)
                        all_embeddings.append(single_embedding)
                        time.sleep(rate_limit_pause)  # 每个请求之间暂停
                    except Exception as inner_e:
                        print(f"处理文本 {j} 失败: {str(inner_e)}")
    
    return all_embeddings

def cosine_similarity(vec1: np.ndarray, vec2: np.ndarray) -> float:
    """
    计算两个向量的余弦相似度
    
    Args:
        vec1: 第一个向量
        vec2: 第二个向量
    
    Returns:
        余弦相似度值，范围为[-1, 1]
        
    Raises:
        ValueError: 如果向量维度不匹配
    """
    if vec1.shape != vec2.shape:
        raise ValueError(f"向量维度不匹配: {vec1.shape} vs {vec2.shape}")
        
    dot_product = np.dot(vec1, vec2)
    norm1 = np.linalg.norm(vec1)
    norm2 = np.linalg.norm(vec2)
    
    # 避免除零错误
    if norm1 == 0 or norm2 == 0:
        return 0.0
    
    return dot_product / (norm1 * norm2)
    

######################### 封装大模型调用函数 #########################

from openai import OpenAI
from zhipuai import ZhipuAI
import os


def call_llm(messages):
    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "sk-8232f1f94ce448f0b5d71563ee23969b"),base_url="https://api.deepseek.com/v1")
    
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=messages,
        temperature=0.7
    )
    
    return response.choices[0].message.content

def call_zhipuai(messages):

    client = ZhipuAI(api_key=os.environ.get("ZHIPUAI_API_KEY", "9e2458a8acf46274d1c1b9418eec500c.oBojpibVHt3IyXnz")) # 填写您自己的APIKey
    response = client.chat.completions.create(
        model="glm-4-flash",  # 填写需要调用的模型编码
        messages=messages,
    )
    # print(response.choices[0].message)
    return response.choices[0].message.content


if __name__ == "__main__":
    # # 测试embedding函数
    # print("测试embedding-2模型:")
    # print(get_zhipuai_embedding("你好"))

    # print("\n测试批量embedding:")
    # print(batch_get_zhipuai_embedding(["你好", "世界"]))

    # # 修正消息格式
    # print("\n测试LLM调用:")
    # test_messages = [
    #     {"role": "user", "content": "你好"}
    # ]
    # try:
    #     print(call_llm(test_messages))
    # except Exception as e:
    #     print(f"调用失败: {e}")

    # print("\n测试智谱AI调用:")
    # test_messages = [
    #     {"role": "user", "content": "你好"}
    # ]
    # try:
    #     print(call_zhipuai(test_messages))
    # except Exception as e:
    #     print(f"调用失败: {e}")

    print("\nMarkdown切片测试:")
    
    # 测试用的Markdown文本
    test_markdown = """# 标题一
这是标题一的内容
一些内容

## 子标题
子标题的内容

# 标题二
这是标题二的内容
"""
    
    print("测试文本:")
    print("-" * 40)
    print(test_markdown)
    print("-" * 40)
    
    print("\n切片结果:")
    chunks = chunk_markdown(test_markdown)
    for i, chunk in enumerate(chunks):
        print(f"\n--- 切片 {i+1} ---")
        print(chunk)
    
    # 测试原始输入
    print("\n测试原始输入:")
    # 非标准格式，没有空格的标题
    original_test = """
# 你好
1111
# 世界
# ## 2你好
# 世界
"""
    print("-" * 40)
    print(original_test)
    print("-" * 40)
    
    print("\n原始输入切片结果:")
    chunks = chunk_markdown(original_test)
    for i, chunk in enumerate(chunks):
        print(f"\n--- 切片 {i+1} ---")
        print(chunk)
        
    # 测试更复杂的情况
    print("\n测试更复杂的情况:")
    complex_test = """# 一级标题
这是一级标题的内容

## 二级标题
二级标题内容

### 三级标题
三级标题内容

# 另一个一级标题
另一个一级标题的内容
"""
    print("-" * 40)
    print(complex_test)
    print("-" * 40)
    
    print("\n复杂情况切片结果:")
    chunks = chunk_markdown(complex_test)
    for i, chunk in enumerate(chunks):
        print(f"\n--- 切片 {i+1} ---")
        print(chunk)


